# -*- coding: utf-8 -*-
# time: 2025/3/29 10:55
# file: ch01.py
# author: hanson
"""
Embedding Models：chromadb 内置的 all-MiniLM-L6-v2

"""
import chromadb.config
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings

# 1. 加载文档（示例使用单个文本文件）
VECTOR_DIR = "data/vector_store"
# 3. 初始化模型 deepseek-r1:7b qwen2.5:7b
embeddings = OllamaEmbeddings(model="qwen2.5:7b")
# 4. 创建向量数据库（自动持久化）


vector_db = Chroma(
    #documents=split_docs,
    persist_directory=VECTOR_DIR,
    embedding_function=embeddings
)
print("向量数据库已创建")

# 5. 构建检索链（优化中文模板）

prompt_template = """
您是一个设计用于査询文档来回答问题的代理。您可以使用文档检索工具。
并基于检索内容来回答问题您可能不查询文档就知道答案，但是您仍然应该查询文档来获得答案。
如果您从文档中找不到任何信息用于回答问题，则只需返回“抱歉，这个问题我还不知道。”作为答案。
如果有人提问等关于您的名字的问题，您就回答：“我是超级牛逼哄哄的小天才助手”作为答案。
上下文：{context}
问题：{question}
"""

retriever = vector_db.as_retriever(
    search_type="mmr",  # 最大边际相关性搜索（平衡相关性与多样性）
    search_kwargs={
        "k": 2,  # 初始检索文档数量
        "fetch_k": 20,  # 底层搜索数量（越大精度越高）
        "lambda_mult": 0.5,  # 多样性控制参数（0-1，越大越多样）
        "score_threshold": 0.3,  # 相关性阈值
    },
)
"""
现实使用过程中相似性查询不是特别准确，我们一般会和全文索引结果结合起来一起评估。
"""
query = "hanson是谁？"
docs = retriever.get_relevant_documents(query)
# 打印结果
for i, doc in enumerate(docs):
    print(f"文档 {i+1}:")
    print(f"内容: {doc.page_content}")
    print(f"来源: {doc.metadata.get('source', '未知')}")
    print("-" * 50)
